{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "de592a7c",
      "metadata": {
        "id": "e4e00c7af679"
      },
      "source": [
        "# Private Fine-Tuning of VaultGemma with LoRA and Differential Privacy\n",
        "\n",
        "This notebook demonstrates how to fine-tune Google's VaultGemma 1B model on medical data using:\n",
        "- **LoRA (Low-Rank Adaptation)**: Efficient parameter-efficient fine-tuning\n",
        "- **4-bit Quantization**: Reduced memory footprint using BitsAndBytes\n",
        "- **Differential Privacy**: Privacy-preserving training with Opacus\n",
        "\n",
        "The goal is to create a medical Q&A model while maintaining strong privacy guarantees.\n",
        "%% [markdown]\n",
        "## 1. Import Libraries and Load Dataset\n",
        "\n",
        "We start by importing all necessary libraries:\n",
        "- `transformers`: For model and tokenizer\n",
        "- `peft`: For LoRA adapters\n",
        "- `opacus`: For differential privacy\n",
        "- `datasets`: For loading and processing the medical dataset\n",
        "\n",
        "The dataset used is **Medical Meadow Medical Flashcards**, which contains medical question-answer pairs.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "745037f8-dd3d-4c04-841a-c1aa13b1ea8c",
      "metadata": {
        "id": "d2c0172b5d9a"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "'\\n!pip install git+https://github.com/huggingface/transformers@v4.56.1-Vault-Gemma-preview\\n! pip install kagglehub\\n! pip install ipywidgets\\n! pip install protobuf -q\\n! pip install tiktoken -q\\n! pip install blobfile -q\\n! pip install sentencepiece -q\\n!pip install -q opacus\\n'"
            ]
          },
          "execution_count": 1,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "# 1. Install necessary libraries\n",
        "!pip install -q -U peft accelerate bitsandbytes datasets pandas\n",
        "!pip install git+https://github.com/huggingface/transformers@v4.56.1-Vault-Gemma-preview\n",
        "!pip install kagglehub ipywidgets opacus -q"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "id": "fadfcc25",
      "metadata": {
        "id": "c05292e28b42"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "input          What is the relationship between very low Mg2+...\n",
            "output         Very low Mg2+ levels correspond to low PTH lev...\n",
            "instruction                      Answer this question truthfully\n",
            "Name: 0, dtype: object\n"
          ]
        }
      ],
      "source": [
        "import os\n",
        "import math\n",
        "import torch\n",
        "import pandas as pd\n",
        "import kagglehub\n",
        "from datasets import load_dataset, Dataset\n",
        "from transformers import (\n",
        "    AutoModelForCausalLM,\n",
        "    GemmaTokenizer,\n",
        "    DataCollatorForLanguageModeling,\n",
        "    BitsAndBytesConfig,\n",
        "    get_cosine_schedule_with_warmup\n",
        ")\n",
        "from peft import LoraConfig, get_peft_model\n",
        "from opacus import PrivacyEngine\n",
        "from opacus.validators import ModuleValidator\n",
        "from torch.utils.data import DataLoader\n",
        "from tqdm.auto import tqdm\n",
        "from peft import PeftModel\n",
        "\n",
        "# Load medical dataset\n",
        "medical_data = load_dataset(\"medalpaca/medical_meadow_medical_flashcards\", split=\"train\")\n",
        "NUM_SAMPLES = 1000\n",
        "data = medical_data.to_pandas().head(NUM_SAMPLES)\n",
        "print(data.iloc[0])"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "fb01bf8d",
      "metadata": {
        "id": "6951338e43f7"
      },
      "source": [
        "## 2. Load Base Model with 4-bit Quantization\n",
        "\n",
        "We load **VaultGemma 1B** from Kaggle with 4-bit quantization to reduce memory usage:\n",
        "- **NF4 quantization**: Normal Float 4-bit quantization for optimal quality\n",
        "- **Double quantization**: Further compression by quantizing the quantization constants\n",
        "- **bfloat16 compute**: Uses brain floating point for stable training\n",
        "\n",
        "The model is automatically distributed across available GPUs using `device_map=\"auto\"`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "id": "fe5002a9",
      "metadata": {
        "id": "14953ef0933d"
      },
      "outputs": [],
      "source": [
        "# Download model from Kaggle\n",
        "model_path = kagglehub.model_download(\"google/vaultgemma/transformers/1b\")\n",
        "\n",
        "# Configure 4-bit quantization\n",
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        ")\n",
        "\n",
        "# Load base model\n",
        "model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_path,\n",
        "    quantization_config=quantization_config,\n",
        "    dtype=torch.bfloat16,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "\n",
        "# Load tokenizer\n",
        "tokenizer = GemmaTokenizer.from_pretrained(model_path)\n",
        "tokenizer.pad_token = tokenizer.eos_token"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "bce6058c",
      "metadata": {
        "id": "43d43a79f42a"
      },
      "source": [
        "## 3. Apply LoRA Adapters\n",
        "\n",
        "**LoRA (Low-Rank Adaptation)** adds small trainable matrices to specific layers while keeping the base model frozen:\n",
        "- **r=8**: Rank of the low-rank matrices (higher = more capacity but more parameters)\n",
        "- **lora_alpha=16**: Scaling factor for LoRA weights\n",
        "- **target_modules**: Which attention and MLP layers to adapt (all projection layers in Gemma)\n",
        "- **lora_dropout=0.05**: Dropout for regularization\n",
        "\n",
        "This approach trains only ~1-2% of the total parameters, making training much faster and memory-efficient.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "id": "e91b3c6f",
      "metadata": {
        "id": "9f332c16c374"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Model and LoRA adapters loaded for training!\n",
            "trainable params: 6,842,368 || all params: 1,045,583,488 || trainable%: 0.6544\n"
          ]
        },
        {
          "data": {
            "text/plain": [
              "PeftModelForCausalLM(\n",
              "  (base_model): LoraModel(\n",
              "    (model): VaultGemmaForCausalLM(\n",
              "      (model): VaultGemmaModel(\n",
              "        (embed_tokens): Embedding(256000, 1152, padding_idx=0)\n",
              "        (layers): ModuleList(\n",
              "          (0-25): 26 x VaultGemmaDecoderLayer(\n",
              "            (self_attn): VaultGemmaAttention(\n",
              "              (q_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (k_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (v_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (o_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1024, out_features=1152, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1024, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1152, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "            )\n",
              "            (mlp): VaultGemmaMLP(\n",
              "              (gate_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=6912, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=6912, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (up_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=6912, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=6912, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (down_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=6912, out_features=1152, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.05, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=6912, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1152, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (act_fn): PytorchGELUTanh()\n",
              "            )\n",
              "            (input_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "            (pre_feedforward_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "          )\n",
              "        )\n",
              "        (norm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "        (rotary_emb): VaultGemmaRotaryEmbedding()\n",
              "      )\n",
              "      (lm_head): Linear(in_features=1152, out_features=256000, bias=False)\n",
              "    )\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 4,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "lora_config = LoraConfig(\n",
        "    r=8,\n",
        "    lora_alpha=16,\n",
        "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
        "    lora_dropout=0.05,\n",
        "    bias=\"none\",\n",
        "    task_type=\"CAUSAL_LM\"\n",
        ")\n",
        "\n",
        "# Apply LoRA adapters to the model (for new training)\n",
        "peft_model = get_peft_model(model, lora_config)\n",
        "\n",
        "print(\"Model and LoRA adapters loaded for training!\")\n",
        "peft_model.print_trainable_parameters()\n",
        "\n",
        "# Set model to training mode\n",
        "peft_model.train()"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "46f3ba58",
      "metadata": {
        "id": "5570b06fb0a5"
      },
      "source": [
        "## 4. Prepare and Tokenize Dataset\n",
        "\n",
        "We create a custom tokenization function that:\n",
        "1. Formats each example as an instruction-following prompt\n",
        "2. Tokenizes the full text (question + answer)\n",
        "3. **Masks the prompt tokens** in the labels by setting them to -100\n",
        "\n",
        "This ensures the model only learns to generate the **response**, not to repeat the question.\n",
        "The masking prevents the loss function from penalizing the model for the input prompt.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "id": "229ae830",
      "metadata": {
        "id": "980e32e29358"
      },
      "outputs": [
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "502ffd5d5ac74087b50623a8b4438814",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "Map:   0%|          | 0/1000 [00:00<?, ? examples/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "def tokenize_and_mask(samples):\n",
        "    \"\"\"\n",
        "    Tokenizes the input and output, then masks the prompt tokens in labels\n",
        "    so that the model only learns to predict the response.\n",
        "    \"\"\"\n",
        "    # Format prompts and responses\n",
        "    full_prompts = [\n",
        "        f\"Instruction:\\nAnswer this question truthfully.\\n\\nQuestion:\\n{inp}\" \n",
        "        for inp in samples[\"input\"]\n",
        "    ]\n",
        "    responses = [f\"\\n\\nResponse:\\n{out}\" for out in samples[\"output\"]]\n",
        "    \n",
        "    # Tokenize full text (prompt + response)\n",
        "    model_inputs = tokenizer(\n",
        "        [p + r for p, r in zip(full_prompts, responses)],\n",
        "        truncation=True,\n",
        "        max_length=128,\n",
        "        padding=\"max_length\",\n",
        "        return_tensors=\"pt\"\n",
        "    )\n",
        "    \n",
        "    # Tokenize only prompts to determine their length\n",
        "    prompt_tokens = tokenizer(\n",
        "        full_prompts,\n",
        "        truncation=True,\n",
        "        max_length=128,\n",
        "        padding=\"max_length\",\n",
        "        return_tensors=\"pt\"\n",
        "    )\n",
        "    \n",
        "    # Create labels (copy of input_ids)\n",
        "    labels = model_inputs[\"input_ids\"].clone()\n",
        "    \n",
        "    # Mask prompt tokens in labels (set to -100 so they're ignored in loss calculation)\n",
        "    for i in range(len(labels)):\n",
        "        prompt_len = int(prompt_tokens[\"attention_mask\"][i].sum())\n",
        "        labels[i][:prompt_len] = -100\n",
        "    \n",
        "    model_inputs[\"labels\"] = labels\n",
        "    return model_inputs\n",
        "\n",
        "# Convert pandas DataFrame to Dataset\n",
        "dataset = Dataset.from_pandas(data)\n",
        "\n",
        "# Apply tokenization function\n",
        "tokenized_dataset = dataset.map(\n",
        "    tokenize_and_mask,\n",
        "    batched=True,\n",
        "    remove_columns=dataset.column_names\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4c6fe678",
      "metadata": {
        "id": "10e3d4fe8c88"
      },
      "source": [
        "## 5. Configure Training Parameters\n",
        "\n",
        "We set up all training hyperparameters and create data loaders:\n",
        "- **90/10 train/validation split** for monitoring overfitting\n",
        "- **Batch size = 1** with **gradient accumulation = 8** (effective batch size of 8)\n",
        "- **Learning rate = 2e-5** with cosine decay schedule\n",
        "- **20 epochs** of training\n",
        "\n",
        "The small batch size is necessary due to memory constraints from the quantized model.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "id": "9f1ee89f",
      "metadata": {
        "id": "a87d38c8c189"
      },
      "outputs": [],
      "source": [
        "# Split dataset into train and validation\n",
        "train_size = int(0.9 * len(tokenized_dataset))\n",
        "train_dataset = tokenized_dataset.select(range(train_size))\n",
        "eval_dataset = tokenized_dataset.select(range(train_size, len(tokenized_dataset)))\n",
        "\n",
        "# Training hyperparameters\n",
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
        "num_train_epochs = 2\n",
        "per_device_train_batch_size = 1\n",
        "gradient_accumulation_steps = 8\n",
        "learning_rate = 2e-5\n",
        "eval_steps = 100\n",
        "logging_steps = 40\n",
        "\n",
        "# Initialize optimizer\n",
        "optimizer = torch.optim.AdamW(peft_model.parameters(), lr=learning_rate)\n",
        "\n",
        "# Data collator for language modeling\n",
        "data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)\n",
        "\n",
        "# Create data loaders\n",
        "train_dataloader = DataLoader(\n",
        "    train_dataset, \n",
        "    batch_size=per_device_train_batch_size, \n",
        "    shuffle=True,\n",
        "    collate_fn=data_collator\n",
        ")\n",
        "eval_dataloader = DataLoader(\n",
        "    eval_dataset, \n",
        "    batch_size=per_device_train_batch_size,\n",
        "    collate_fn=data_collator\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "4710f3de",
      "metadata": {
        "id": "f5282c2b73bc"
      },
      "source": [
        "## 6. Apply Differential Privacy with Opacus\n",
        "\n",
        "**Differential Privacy (DP)** ensures that the trained model doesn't memorize specific training examples:\n",
        "- **ε (epsilon) = 8.0**: Privacy budget (lower = more privacy, but potentially worse performance)\n",
        "- **δ (delta) = 1e-5**: Probability of privacy breach (should be < 1/dataset_size)\n",
        "- **max_grad_norm = 1.0**: Clips gradients to prevent any single example from having too much influence\n",
        "\n",
        "Opacus modifies the training loop to add calibrated noise to gradients, providing mathematical privacy guarantees.\n",
        "The final epsilon value will tell us exactly how much privacy was consumed during training.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "bf0b24b8",
      "metadata": {
        "id": "81dc25216480"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/home/neo/anaconda3/envs/torch/lib/python3.11/site-packages/opacus/privacy_engine.py:96: UserWarning: Secure RNG turned off. This is perfectly fine for experimentation as it allows for much faster training performance, but remember to turn it on and retrain one last time before production with ``secure_mode`` turned on.\n",
            "  warnings.warn(\n",
            "/home/neo/anaconda3/envs/torch/lib/python3.11/site-packages/opacus/accountants/analysis/rdp.py:332: UserWarning: Optimal order is the largest alpha. Please consider expanding the range of alphas to get a tighter privacy bound.\n",
            "  warnings.warn(\n",
            "/home/neo/anaconda3/envs/torch/lib/python3.11/site-packages/opacus/accountants/analysis/rdp.py:332: UserWarning: Optimal order is the smallest alpha. Please consider expanding the range of alphas to get a tighter privacy bound.\n",
            "  warnings.warn(\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using cosine learning rate schedule with warmup.\n"
          ]
        }
      ],
      "source": [
        "# Differential privacy setup\n",
        "target_delta = 1e-5   # Lower values = more privacy\n",
        "target_epsilon = 15.0 # Lower values = more privacy\n",
        "\n",
        "privacy_engine = PrivacyEngine()\n",
        "peft_model, optimizer, train_dataloader = privacy_engine.make_private_with_epsilon(\n",
        "    module=peft_model,\n",
        "    optimizer=optimizer,\n",
        "    data_loader=train_dataloader,\n",
        "    target_epsilon=target_epsilon,\n",
        "    target_delta=target_delta,\n",
        "    epochs=num_train_epochs,\n",
        "    max_grad_norm=1.0,\n",
        "    poisson_sampling=False\n",
        ")\n",
        "\n",
        "# Validate and fix model for Opacus compatibility\n",
        "if not ModuleValidator.is_valid(peft_model):\n",
        "    peft_model = ModuleValidator.fix(peft_model)\n",
        "\n",
        "peft_model.train()\n",
        "peft_model.to(device)\n",
        "\n",
        "# Learning rate scheduler with cosine decay and warmup\n",
        "num_training_steps = math.ceil(len(train_dataloader) / gradient_accumulation_steps) * num_train_epochs\n",
        "num_warmup_steps = 40\n",
        "\n",
        "lr_scheduler = get_cosine_schedule_with_warmup(\n",
        "    optimizer=optimizer,\n",
        "    num_warmup_steps=num_warmup_steps,\n",
        "    num_training_steps=num_training_steps\n",
        ")\n",
        "\n",
        "print(\"Using cosine learning rate schedule with warmup.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "b84fda3c",
      "metadata": {
        "id": "9bd49b322ded"
      },
      "source": [
        "## 8. Training Loop\n",
        "\n",
        "The main training loop with the following features:\n",
        "- **Gradient accumulation**: Accumulates gradients over 8 steps before updating weights\n",
        "- **Automatic checkpointing**: Saves model when training loss drops below 0.06\n",
        "- **Periodic validation**: Evaluates on validation set every 200 steps\n",
        "- **Progress tracking**: Uses tqdm for visual progress bar\n",
        "\n",
        "The loop will run for 2 epochs, logging metrics every 20 steps and evaluating every 200 steps.\n",
        "Models are saved to the specified directory when performance thresholds are met.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "05f7e60d",
      "metadata": {
        "id": "087239f697b1"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Starting training loop...\n"
          ]
        },
        {
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "44556abde9714e90b05c72706198bb85",
              "version_major": 2,
              "version_minor": 0
            },
            "text/plain": [
              "  0%|          | 0/226 [00:00<?, ?it/s]"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "sys:1: UserWarning: Full backward hook is firing when gradients are computed with respect to module outputs since no inputs require gradients. See https://docs.pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook for more details.\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Step 40: Train Loss = 23.2601\n",
            "Step 80: Train Loss = 22.7426\n",
            "Step 120: Train Loss = 4.4020\n",
            "Step 160: Train Loss = 22.0963\n",
            "Step 200: Train Loss = 21.9594 | Validation Loss = 2.8816\n",
            "Final privacy cost: ε = 22.21 for δ = 0.01\n"
          ]
        }
      ],
      "source": [
        "print(\"Starting training loop...\")\n",
        "progress_bar = tqdm(range(num_training_steps))\n",
        "global_step = 0\n",
        "\n",
        "for epoch in range(num_train_epochs):\n",
        "    peft_model.train()\n",
        "    train_loss_accumulator = 0.0\n",
        "    \n",
        "    for step, batch in enumerate(train_dataloader):\n",
        "        # Move batch to device\n",
        "        batch = {k: v.to(device) for k, v in batch.items()}\n",
        "        \n",
        "        # Forward pass\n",
        "        outputs = peft_model(**batch)\n",
        "        loss = outputs.loss\n",
        "        train_loss_accumulator += loss.item()\n",
        "        \n",
        "        # Backward pass\n",
        "        loss.backward()\n",
        "        \n",
        "        # Optimizer step with gradient accumulation\n",
        "        if (step + 1) % gradient_accumulation_steps == 0:\n",
        "            optimizer.step()\n",
        "            lr_scheduler.step()\n",
        "            optimizer.zero_grad()\n",
        "            \n",
        "            global_step += 1\n",
        "            progress_bar.update(1)\n",
        "            \n",
        "            # Logging and checkpoint saving\n",
        "            if global_step % logging_steps == 0:\n",
        "                avg_train_loss = train_loss_accumulator / logging_steps\n",
        "                log_message = f\"Step {global_step}: Train Loss = {avg_train_loss:.4f}\"\n",
        "                \n",
        "                # Save checkpoint if loss is below threshold\n",
        "                if avg_train_loss < 0.06:\n",
        "                    checkpoint_path = \"./final_model\"\n",
        "                    \n",
        "                    # Save PEFT adapters and tokenizer\n",
        "                    peft_model.save_pretrained(checkpoint_path)\n",
        "                    tokenizer.save_pretrained(checkpoint_path)\n",
        "                    log_message += f\" | Model Saved to {checkpoint_path}\"\n",
        "                \n",
        "                # Validation evaluation\n",
        "                if global_step % eval_steps == 0:\n",
        "                    peft_model.eval()\n",
        "                    eval_losses = []\n",
        "                    \n",
        "                    with torch.no_grad():\n",
        "                        for eval_batch in eval_dataloader:\n",
        "                            eval_batch = {k: v.to(device) for k, v in eval_batch.items()}\n",
        "                            eval_outputs = peft_model(**eval_batch)\n",
        "                            eval_losses.append(eval_outputs.loss.item())\n",
        "                    \n",
        "                    avg_eval_loss = sum(eval_losses) / len(eval_losses)\n",
        "                    log_message += f\" | Validation Loss = {avg_eval_loss:.4f}\"\n",
        "                    peft_model.train()\n",
        "                \n",
        "                print(log_message)\n",
        "                train_loss_accumulator = 0.0\n",
        "\n",
        "# Final privacy budget\n",
        "epsilon = privacy_engine.get_epsilon(delta=target_delta)\n",
        "print(f\"Final privacy cost: ε = {epsilon:.2f} for δ = {target_delta}\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "f127c5dd",
      "metadata": {
        "id": "79a4047ea492"
      },
      "outputs": [
        {
          "data": {
            "text/plain": [
              "PeftModelForCausalLM(\n",
              "  (base_model): LoraModel(\n",
              "    (model): VaultGemmaForCausalLM(\n",
              "      (model): VaultGemmaModel(\n",
              "        (embed_tokens): Embedding(256000, 1152, padding_idx=0)\n",
              "        (layers): ModuleList(\n",
              "          (0-25): 26 x VaultGemmaDecoderLayer(\n",
              "            (self_attn): VaultGemmaAttention(\n",
              "              (q_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (k_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (v_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=1024, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1024, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (o_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1024, out_features=1152, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1024, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1152, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "            )\n",
              "            (mlp): VaultGemmaMLP(\n",
              "              (gate_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=6912, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=6912, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (up_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=1152, out_features=6912, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=1152, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=6912, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (down_proj): lora.Linear4bit(\n",
              "                (base_layer): Linear4bit(in_features=6912, out_features=1152, bias=False)\n",
              "                (lora_dropout): ModuleDict(\n",
              "                  (default): Dropout(p=0.1, inplace=False)\n",
              "                )\n",
              "                (lora_A): ModuleDict(\n",
              "                  (default): Linear(in_features=6912, out_features=8, bias=False)\n",
              "                )\n",
              "                (lora_B): ModuleDict(\n",
              "                  (default): Linear(in_features=8, out_features=1152, bias=False)\n",
              "                )\n",
              "                (lora_embedding_A): ParameterDict()\n",
              "                (lora_embedding_B): ParameterDict()\n",
              "                (lora_magnitude_vector): ModuleDict()\n",
              "              )\n",
              "              (act_fn): PytorchGELUTanh()\n",
              "            )\n",
              "            (input_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "            (pre_feedforward_layernorm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "          )\n",
              "        )\n",
              "        (norm): VaultGemmaRMSNorm((1152,), eps=1e-06)\n",
              "        (rotary_emb): VaultGemmaRotaryEmbedding()\n",
              "      )\n",
              "      (lm_head): Linear(in_features=1152, out_features=256000, bias=False)\n",
              "    )\n",
              "  )\n",
              ")"
            ]
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ],
      "source": [
        "model_path = kagglehub.model_download(\"google/vaultgemma/transformers/1b\")\n",
        "\n",
        "quantization_config = BitsAndBytesConfig(\n",
        "    load_in_4bit=True,\n",
        "    bnb_4bit_quant_type=\"nf4\",\n",
        "    bnb_4bit_compute_dtype=torch.bfloat16,\n",
        "    bnb_4bit_use_double_quant=True,\n",
        ")\n",
        "\n",
        "base_model = AutoModelForCausalLM.from_pretrained(\n",
        "    model_path,\n",
        "    quantization_config=quantization_config,\n",
        "    dtype=torch.bfloat16,\n",
        "    device_map=\"auto\",\n",
        ")\n",
        "\n",
        "adapter_path = \"./final_model\"\n",
        "\n",
        "tokenizer = GemmaTokenizer.from_pretrained(adapter_path)\n",
        "tokenizer.pad_token = tokenizer.eos_token\n",
        "\n",
        "peft_model = PeftModel.from_pretrained(base_model, adapter_path, is_trainable=False)\n",
        "\n",
        "peft_model.eval()\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4a2abf81",
      "metadata": {
        "id": "8a4747527d13"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "\n",
            "Question: What is the role of insulin in the human body?\n",
            "Answer: The hormone affects several important functions including its actions on blood glucose levels which control hormones to lower or maintain it so that we can feel more adequate and are able respond better\n"
          ]
        }
      ],
      "source": [
        "def generate_response(question, max_new_tokens=128, temperature=0.9, top_p=0.9):\n",
        "    prompt = f\"Instruction:\\nAnswer this question truthfully.\\n\\nQuestion:\\n{question}\\n\\nResponse:\\n\"\n",
        "    \n",
        "    inputs = tokenizer(prompt, return_tensors=\"pt\", padding=True)\n",
        "    inputs = {k: v.to(peft_model.device) for k, v in inputs.items()}\n",
        "    \n",
        "    with torch.no_grad():\n",
        "        outputs = peft_model.generate(\n",
        "            **inputs,\n",
        "            max_new_tokens=max_new_tokens,\n",
        "            temperature=temperature,\n",
        "            top_p=top_p,\n",
        "            do_sample=True,\n",
        "            eos_token_id=tokenizer.eos_token_id,\n",
        "            repetition_penalty=1.4 \n",
        "        )\n",
        "    \n",
        "    full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
        "    \n",
        "    if \"Response:\" in full_response:\n",
        "        response = full_response.split(\"Response:\")[-1].strip()\n",
        "    else:\n",
        "        response = full_response\n",
        "    \n",
        "    return response\n",
        "\n",
        "\n",
        "question = \"What is the role of insulin in the human body?\"\n",
        "\n",
        "response = generate_response(question)\n",
        "print(f\"\\nQuestion: {question}\")\n",
        "print(f\"Answer: {response}\")\n"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "[VaultGemma]FineTuning_Inference_Huggingface.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
